"""
    @author: caitao
    @Date: 2024-11-30
    @feature:
        * 通过本地数据库取数，进行黄金台识别
    @update:
        * 修改匹配方法，先过滤涨停日期行索引，再判断后续规则
"""
import pandas as pd
import math
import numpy as np
import sqlalchemy
from sqlalchemy import and_

from utilities.utilities_func import MyThread
from data_resource.data_bases import StockInfo, TsMarketDaily, Session, engine, engine2

num_threads = 20
session = Session()


def stockIncrease_type(code):
    """
    根据股票代码判断涨跌停办要求
    :param code:
    :return:
    """
    if code[0] == '3' or code[0:3] == '688':
        return 'radical'
    else:
        return 'normal'


def get_stocks():
    """
    return: list of tuple('ticker'，)
    """
    return pd.read_sql(session.query(StockInfo.ticker).distinct().filter(StockInfo.status == 'L').statement,
                       session.bind)


def golden_filter(price, start_window, amplitude=15.00, boxrange_bottom=-0.01, boxrange_upper=0.8, v=0.6):
    if price.shape[0] < start_window:
        return False
    # 振幅合集
    amplitude_list = price.iloc[1:, 11].to_list()
    price['vollage'] = abs(price['open'] - price['close'])  # 计算柱体大小

    if all(_ <= amplitude for _ in amplitude_list):
        # 确定箱体上下界
        _close1 = price.iloc[0, 3]
        _open1 = price.iloc[0, 2]
        _v1 = price.iloc[0, -1]
        # 2-窗口末尾 K线柱的开盘价、收盘价合集
        a = price.iloc[1:, 2].to_list() + price.iloc[1:, 3].to_list()
        # 箱体上下限
        bottom = _v1 * boxrange_bottom + _open1
        upper = _v1 * boxrange_upper + _close1
        # 柱体大小限制
        vollage = _v1 * v
        # 2-10K线柱体大小
        b = price.iloc[1:, -1].to_list()
        if (max(a) <= upper) and (min(a) >= bottom) and (max(b) <= vollage) and (_close1 > _open1):
            return True
        else:
            return False
    else:
        return False


def trigger_v2(code, start_date, end_date, main_bottom=9.97, main_upper=11.00, star_bottom=19.97,
               star_upper=21.00,
               amplitude=15.00, boxrange_bottom=-0.01, boxrange_upper=0.8, v=0.6,
               tao=0.5, start_window=5, end_window=11):
    start_date = pd.to_datetime(start_date).date().strftime('%Y-%m-%d')
    end_date = pd.to_datetime(end_date).date().strftime('%Y-%m-%d')

    # 获取历史数据
    _stockprice = pd.read_sql(
        session.query(TsMarketDaily).filter(
            and_(TsMarketDaily.ticker == code, start_date <= TsMarketDaily.trade_date,
                 TsMarketDaily.trade_date <= end_date)
        ).order_by(TsMarketDaily.trade_date).statement,
        session.bind
    )

    _stockprice.drop(columns=['id'], inplace=True)
    _stockprice.rename(columns={'trade_date': 'tradeDate'}, inplace=True)
    _stockprice.reset_index(drop=True, inplace=True)

    _stockprice['tradeDate'] = pd.to_datetime(_stockprice['tradeDate'])
    _stockprice['tradeDate'] = _stockprice['tradeDate'].dt.date
    # 计算股票振幅
    _stockprice['amplitude'] = (_stockprice['high'] - _stockprice['low']) / _stockprice['pre_close'] * 100

    # 过滤涨停
    _temp = stockIncrease_type(code)

    def _core(_stockprice, _temp):
        _results = []
        if _temp == 'normal':
            _raw = _stockprice[
                (_stockprice['pct_chg'] > main_bottom) & (_stockprice['pct_chg'] < main_upper)].index
        else:
            _raw = _stockprice[
                (_stockprice['pct_chg'] > star_bottom) & (_stockprice['pct_chg'] < star_upper)].index
        if _raw.empty:
            return _results
        else:
            for i in _raw:
                # 定义滑动窗口
                for w in range(int(start_window), int(end_window), 1):
                    _e = i + w
                    _price = _stockprice.iloc[i:_e, :].copy()

                    if len(_price) < start_window:
                        continue
                    else:
                        # 黄金台2-10柱体逻辑判断，固定窗口
                        _x = golden_filter(_price, start_window, amplitude, boxrange_bottom, boxrange_upper, v)
                        # 是否收敛判断
                        _upper = _price['high'].to_list()[1:]
                        _low = _price.iloc[1:, 6].to_list()
                        _distance = [a - b for a, b in zip(_upper, _low)]
                        if all(item == 0 for item in _distance):
                            _c = True
                        elif np.mean(_distance) < tao:
                            _c = True
                        else:
                            _c = False
                        # 判断条件：黄金台或者极致收敛任意触发
                        if _x or _c:
                            _results.append({'ticker': code, 'signal': _price.iloc[-1, 0]})
                            # print(f"---黄金台信号---{code}, 信号日：{_price.iloc[_e, 0]}")
                            continue
                        else:
                            continue
            return _results

    _result = _core(_stockprice, _temp)
    return _result


def goldenFilter_pool(codelist, start_date, end_date,
                      main_bottom=9.97, main_upper=11.00, star_bottom=19.97, star_upper=21.00,
                      amplitude=15.00, boxrange_bottom=-0.01, boxrange_upper=0.8, v=0.6,
                      tao=0.5, start_window=5, end_window=11):
    results = []
    for i in codelist:
        _a = trigger_v2(i, start_date, end_date, main_bottom, main_upper, star_bottom, star_upper,
                        amplitude, boxrange_bottom, boxrange_upper, v,
                        tao, start_window, end_window)
        if len(_a) > 0:
            results.extend(_a)
        else:
            continue
    return results


def goldenFilterThreads(start_date, end_date,
                        main_bottom=9.97, main_upper=11.00, star_bottom=19.97, star_upper=21.00,
                        amplitude=15.00, boxrange_bottom=-0.01, boxrange_upper=0.8, v=0.6,
                        tao=0.5, start_window=5, end_window=11):

    stock = get_stocks()
    results = []
    f = locals()
    for thread in range(num_threads):
        delta = math.ceil(stock.shape[0] / num_threads)
        _stocklist = stock['ticker'].to_list()[thread * delta: (thread + 1) * delta]
        f['n_' + str(thread)] = MyThread(goldenFilter_pool, args=(_stocklist, start_date, end_date,
                                                                  main_bottom, main_upper, star_bottom, star_upper,
                                                                  amplitude, boxrange_bottom, boxrange_upper, v,
                                                                  tao, start_window, end_window
                                                                  ))

    for t in range(num_threads):
        f['n_' + str(t)].start()

    for t in range(num_threads):
        f['n_' + str(t)].join()

    # 获取结果
    for t in range(num_threads):
        _result = f['n_' + str(t)].get_result()
        results.extend(_result)

    return results


def golden_daily(start_date, end_date):
    x = goldenFilterThreads(start_date, end_date)
    result = []
    for i in x:
        _ = pd.DataFrame.from_dict(i, orient='index').T
        result.append(_)
    result = pd.concat(result)
    result.reset_index(drop=True, inplace=True)
    result.rename(columns={'signal': 'signal_date'}, inplace=True)

    d_type1 = {
        'ticker': sqlalchemy.types.VARCHAR, 'signal': sqlalchemy.types.DATE,
    }

    num = 0
    print("-----------开始写入黄金台计算结果-------------------")
    for m in range(len(result)):
        _x = pd.DataFrame(result.iloc[m, :]).T
        try:
            _x.to_sql('factor_golden', con=engine, if_exists='append',
                      index=False, schema='quant_research',
                      dtype=d_type1, chunksize=100000)

            _x.to_sql(
                'results_golden', engine2, if_exists='append',
                index=False, dtype=d_type1
            )

            num += 1
        except Exception as e:
            continue
    print(f"--------------------黄金台结果写入完成，共写入{num}条数据---------------------")

    return result

